//////////////////////////////////////////////////////////////////
// MattiseOS Kernel
// Date: 2008
// Author: Matthew Iselin
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// at your option) any later version.
// 
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
// 
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.
//////////////////////////////////////////////////////////////////

// Includes
#include <mattise.h>
#include <stdint.h>
#include <mem/util.h>
#include <mem/heap.h>
#include <mem/page.h>
#include <kernio/io.h>
#include <process/common.h>

// Functions

// I'm going to use a "heap" method for kernel stacks, that will take up a percentage of system RAM.
// For instance, in a 256 MB system it will take ~ 25MB, which is 10%. Each stack is around 2 pages,
// which is 8 KB. Therefore, in 25 MB you can have over 3,000 individual kernel threads.
//
// Stacks are stored starting at 0x50000000, giving userspace tasks room for a 1 GB heap.

// Types 
// a block of memory in the address space
struct memblock
{
	uint32_t used; // 4 bytes
	uint32_t size; // 4 bytes
	struct memblock* next; // 4 bytes
} __attribute__((packed)); // 12 bytes

// Variables

// the stack heap
static struct memblock* stack_heap = 0;

// total available RAM
static uint32_t total_avail_ram = 0;

// the current kernel stack heap base
uint32_t kstk_heap_base = 0;

// the size of the kernel stack heap initially
#define SZ_KERNEL_STACK_HEAP 1

// Functions

// sets up the kernel stack allocator
void setup_kernel_stack_heap( uint32_t maxmem )
{
	// get the total RAM
	total_avail_ram = maxmem;

	// allocate the space
	kstk_heap_base = 0x50000000;
	stack_heap = (struct memblock*) kAllocateAndMap(SZ_KERNEL_STACK_HEAP, 0x50000000, PF_PRESENT | PF_KERNELPAGE | PF_USER | PF_READWRITE | PF_IMMUTABLE);
	kstk_heap_base += SZ_KERNEL_STACK_HEAP * 0x1000;

	// setup it's state
	stack_heap->used = 0;
	stack_heap->size = SZ_KERNEL_STACK_HEAP * 0x1000;
	stack_heap->next = 0;
}

// allocates memory on the heap
void* get_kernel_stack()
{
	// locals that will be used in the allocation
	struct memblock* mblock, *newblock, *prev;
	char* cblock;
	size_t size2;
	size_t size = SZ_KSTACK * 0x1000;

	// calculate the actual size that the block will require
	size2 = size + sizeof( struct memblock );

	// check that there is already a heap
	if( stack_heap == 0 )
	{
		// find out the total memory we have
		// set the kernel heap location
		dprintf("== [FATAL] No stack heap! ==\n");
		return 0;
	}

	// look for a block (effectively this provides us with a linked list
	// of blocks with each entry being a header to a block of memory)
	for( mblock = stack_heap; mblock != 0; prev = mblock, mblock = mblock->next )
		if( (!mblock->used) && (size2 <= mblock->size) )
			break;

	// check that one was found
	if( mblock == 0 )
	{
		// check for a valid prev variable (needed to work properly)
		if(prev != 0)
		{
			// attempt to extend the heap by a page (or more, depending on size)
			uint32_t extend_sz = size > 0x1000 ? ((size & ~0xFFF) / 0x1000) + 1 : 1;
			uint32_t old_base = kstk_heap_base;
			kstk_heap_base += extend_sz * 0x1000;
			struct memblock* b = (struct memblock*) kAllocateAndMap(extend_sz, old_base, PF_PRESENT | PF_KERNELPAGE | PF_USER | PF_READWRITE | PF_IMMUTABLE);
			if(b == 0)
				mblock = 0;
			else
			{
				// link it in (merge if needed)
				if(prev->used == 0)
					prev->size += extend_sz * 0x1000;
				else
				{
					prev->next = b;
					b->next = 0;
					b->size = extend_sz * 0x1000;
					b->used = 0;
				}

				// search again
				for( mblock = stack_heap; mblock != 0; prev = mblock, mblock = mblock->next )
					if( (!mblock->used) && (size2 <= mblock->size) )
						break;
			}
		}

		// if there still isn't a valid block, give up
		if(mblock == 0)
		{
			dprintf("[MEM  ] kernel stack heap allocator returns null\n");
			return 0;
		}
	}

	// setup the block (cblock holds the entire block)
	cblock = (char*) mblock + size2;
	newblock = (struct memblock*) cblock;

	// point this entry to a new block with our settings
	newblock->used = 0;
	newblock->size = mblock->size - size2;
	newblock->next = mblock->next;

	// point the old block to the new block and set it up
	mblock->used = 1;
	mblock->size = size;
	mblock->next = newblock;

	// return a pointer to the memory
	cblock = (char*) mblock;
	return (void*) (cblock + sizeof( struct memblock ));
}

// frees memory from the heap
void free_kernel_stack( void* ptr )
{
	// locals for use later
	struct memblock* mblock;
	char* cblock;

	// get the real address of the block to free
	cblock = (char*) ptr - sizeof( struct memblock );

	// find this address in the linked list
	for( mblock = stack_heap; mblock != 0; mblock = mblock->next )
	{
		// check that it's equal
		if( mblock == (struct memblock*) cblock )
			break;
	}

	// check that the pointer is a valid entry
	if( mblock == 0 )
		return;

	// free this block
	mblock->used = 0;

	// try to merge two unused blocks together
	for( mblock = stack_heap; mblock != 0; mblock = mblock->next )
	{
		// look for valid blocks to merge
		while( ( !mblock->used ) && ( mblock->next != 0 ) && ( !mblock->next->used ) )
		{
			// merge the two blocks (their size should add up and the pointers should change)
			mblock->size = ((int32_t) mblock->next) - ((uint32_t) mblock) + mblock->next->size;
			mblock->next = mblock->next->next;
		}
	}
}

// sets up the stack bitmaps for a context
void SetUpContextBitmap(struct context* ctx)
{
	// allocate space first
	ctx->kernel_stack_bitmap = 0;
	ctx->user_stack_bitmap = (uint32_t*) malloc( SZ_BITMAP(SZ_PSTACK) );

	// zero them out
	memset( ctx->user_stack_bitmap, 0, SZ_BITMAP(SZ_PSTACK) );
}

// returns a stack to a child thread
void ReturnChildKernelStack(uint32_t stack)
{
	// free it
	free_kernel_stack((void*) stack);
}

// returns a user stack for a child thread
void ReturnChildUserStack(struct context* ctx, uint32_t stack)
{
	// get the parent context
	struct context* parent = ctx->parent;

	// the stack comes in as an address - please, PLEASE, let the compiler
	// optomise this!
	stack -= STACK_USER_VIRT; // subtract the virtual base
	stack /= (0x1000 * SZ_PSTACK); // convert to a page number

	// stack now contains the bit number, i think
	uint32_t offset = stack / 32;
	uint32_t bitnum = stack % 32; // i think?

	parent->user_stack_bitmap[offset] &= ~( 1 << ( bitnum % 32 ) );
}

// gets a stack to use for a new child thread
uint32_t GetChildKernelStack()
{
	return (uint32_t) get_kernel_stack();
}

// gets a stack to use for a new child thread
uint32_t GetChildUserStack(struct context* ctx)
{
	// get the parent context
	struct context* parent = ctx->parent;

	// search in the bitmap
	size_t i;
	for(i = 0; i < SZ_BITMAP(SZ_PSTACK); i++)
	{
		if(parent->user_stack_bitmap[i] == 0xffffffff)
			continue;
		else
		{
			// go through each bit
			size_t j;
			for(j = 0; j < 32; j++)
			{
				if(parent->user_stack_bitmap[i] & ( 1 << ( j % 32 ) ))
					continue;
				parent->user_stack_bitmap[i] |= ( 1 << ( j % 32 ) );
				uint32_t addr = (i * (32 * SZ_PSTACK * 0x1000)) + (j * SZ_PSTACK * 0x1000);
				return addr;
			}
		}
	}
	dprintf("==== [ERROR] no valid stack found ====\n");
	return 0;
}
